import gradio as gr
from RAG import chunking, embedding, vectordb, retrive, rerank, generation
import os


def rag_pipeline(
    query: str,
    data_path: str = "./data/data.json",
    db_path: str = "./vector_db",
    top_k_retrive=20,
    top_k_rerank=10,
):
    """
    RAG完整流程函数

    Args:
        query (str): 用户查询
        data_path (str): 数据文件路径
        db_path (str): 向量数据库路径
        top_k_retrive = 6
        top_k_rerank = 4

    Returns:
        str: 生成的回答
    """


    # 如果数据库不存在，则初始化
    if not os.path.exists(db_path) or not os.listdir(db_path):
        chunks = chunking.chunking(data_path)
        data_embedding = embedding.embedding(chunks)
        vectordb.store_to_db(db_path, data_embedding, chunks)

    # 执行RAG流程
    data_retrive = retrive.retrive(query, db_path, top_k_retrive)
    data_rerank = rerank.rerank(query, data_retrive, top_k_rerank)
    result = generation.generate(query, data_rerank)

    return result


def launch_gradio_interface():
    """启动Gradio界面"""
    with gr.Blocks(title="XDUMsgBot - 西安电子科技大学问答助手") as demo:
        gr.Markdown("# 🎓 XDUMsgBot - 西安电子科技大学问答助手")
        gr.Markdown("请输入关于西安电子科技大学的问题，我会尽力为您解答。")

        with gr.Row():
            with gr.Column():
                query_input = gr.Textbox(
                    label="您的问题",
                    placeholder="例如：学校的奖学金政策是什么？",
                    lines=2,
                )
                submit_btn = gr.Button("获取答案", variant="primary")

            with gr.Column():
                answer_output = gr.Markdown(label="答案")

        # 示例问题
        gr.Examples(
            examples=[
                "学校的奖学金政策是什么？",
                "如何申请研究生？",
                "本科生的培养方案有哪些？",
                "学校有哪些特色专业？",
                "校园生活怎么样？",
            ],
            inputs=query_input,
        )

        # 绑定事件
        submit_btn.click(fn=rag_pipeline, inputs=query_input, outputs=answer_output)

        query_input.submit(fn=rag_pipeline, inputs=query_input, outputs=answer_output)

    return demo


if __name__ == "__main__":
    demo = launch_gradio_interface()
    demo.launch(server_name="localhost", server_port=7860, share=False)